package com.jf.ai.controller;

import com.jf.ai.service.ImageService;
import jakarta.annotation.Resource;
import org.springframework.core.io.InputStreamResource;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.client.RestTemplate;

import java.io.ByteArrayInputStream;

/**
 * Created with IntelliJ IDEA
 * Description:
 * User:22954
 * Date:2025-06-16
 * Time:14:50
 */
@RestController
public class ImageController {
    @Resource
    private ImageService imageService;

    private final RestTemplate restTemplate = new RestTemplate();

    @GetMapping(value = "/generateImage", produces = MediaType.IMAGE_PNG_VALUE)
    public ResponseEntity<InputStreamResource> generateImage(@RequestParam("message") String prompt) {
        String imageUrl = imageService.generateImage2(prompt);

        byte[] imageBytes = restTemplate.getForObject(imageUrl, byte[].class);

        assert imageBytes != null;
        InputStreamResource imageStream = new InputStreamResource(new ByteArrayInputStream(imageBytes));
        return ResponseEntity.ok()
                .header(HttpHeaders.CONTENT_TYPE, MediaType.IMAGE_PNG_VALUE)
                .body(new InputStreamResource(imageStream));
    }
}
